#%%
%matplotlib inline
import numpy as np, pandas as pd, pathlib, matplotlib.pyplot as plt, seaborn as sns
%config InlineBackend.figure_format='retina'
plt.style.use('seaborn-whitegrid')
sns.set_context('paper')
sns.set_style('ticks')
cm = 1/2.54

df = {}
for fn in list(pathlib.Path('4_data_settling_velocities/').glob('*_settling_velocities.csv')):
    df_temp = pd.read_csv(fn)
    df[fn.stem.split('_')[0]] = {'vel_mean': df_temp.vel.mean()/1e3, 'dvel_mean': df_temp.vel.std()/1e3/np.sqrt(df_temp.vel.size), 'vels': df_temp.vel.values/1e3, 'dvels': df_temp.vel_e.values/1e3}

df = pd.DataFrame(df).T

df_vol = pd.read_csv('1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)


df['CD'] = np.nan
for grain in df.vel_mean.index:
    df['CD'][grain] = 4*(df_vol.total_grain_ro[grain]/990 - 1)*9.81*df_shape.dn[grain] / (3*df.vel_mean[grain]**2) * df_shape.CSF[grain]

df.to_csv('4_var_settling_velocities.csv')

df = df.T
fs = 7
fig, axs = plt.subplots(1, 5, figsize=(18*cm, 5*cm), sharey=True)
label = {'gb': 'Spheres', 'oc': 'Faceted ellipsoids', 'gg': 'Rounded chips', 'ls': 'Rectangular prisms', 'ng': 'Natural gravel'}
ns = np.zeros(5)
for i, grain in enumerate(['gb', 'oc', 'gg', 'ng', 'ls']):
        axs[i].hist(df[grain]['vels'], density=True, alpha=.6, bins=20, range=(0,.5), color='k', linewidth=0) # A
        ns[i] = df[grain]['vels'].size
        axs[i].set_title(label[grain] + '\n(n = %i)' % ns[i], fontsize=fs)
        axs[i].tick_params(axis='x', labelsize=fs)
        axs[i].tick_params(axis='y', labelsize=fs)

for i in range(5):
        axs[i].set_xlim(0.2,0.6)
       
axs[2].set_xlabel('Settling velocity (m/s)', fontsize=fs)
axs[0].set_ylabel('# observations', fontsize=fs)

plt.tight_layout()
plt.savefig('Deal_EDfig7.jpg', dpi=300)

# %%
